from typing import Literal, TypedDict, List
from langchain.tools import tool
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import StateGraph


# 1. 扩展状态：新增“是否需要搜索”的标记
class AdvancedChatState(TypedDict):
    messages: List[HumanMessage | AIMessage | str]  # 增加str类型支持系统消息
    need_search: bool  # 新增字段：标记是否需要搜索


# 2. 定义模拟搜索工具
@tool
def search_tool(query: str) -> str:
    """模拟搜索工具，返回查询结果"""
    return f"搜索到关于「{query}」的信息：LangGraph 是 LangChain 生态的工作流框架，支持图结构流程。"


# 1. 判断节点：判断是否需要搜索
def judge_search_node(state: AdvancedChatState) -> AdvancedChatState:
    """根据用户消息，判断是否需要调用搜索工具"""
    user_msg = state["messages"][-1].content if isinstance(state["messages"][-1], (HumanMessage, AIMessage)) else \
    state["messages"][-1]
    # 简单规则：若消息含“是什么”“定义”，则需要搜索（模拟逻辑）
    need_search = "是什么" in user_msg or "定义" in user_msg
    return {"need_search": need_search}  # 仅更新 need_search 字段


# 2. 搜索节点：调用搜索工具，更新状态
def search_node(state: AdvancedChatState) -> AdvancedChatState:
    user_msg = state["messages"][-1].content if isinstance(state["messages"][-1], (HumanMessage, AIMessage)) else \
    state["messages"][-1]
    # 调用搜索工具
    search_result = search_tool.invoke(user_msg)
    # 将搜索结果添加到对话历史（作为“系统消息”）
    system_msg = f"【搜索结果】：{search_result}"
    return {
        "messages": state["messages"] + [system_msg]  # 新增系统消息
    }


from langchain.chat_models import init_chat_model

llm = init_chat_model(model="deepseek-chat", model_provider="deepseek")


# 3. LLM节点：结合状态（含搜索结果，若有）生成回答
def advanced_llm_node(state: AdvancedChatState) -> AdvancedChatState:
    messages = state["messages"]
    # 转换消息格式为LLM可处理的类型
    processed_messages = []
    for msg in messages:
        if isinstance(msg, (HumanMessage, AIMessage)):
            processed_messages.append(msg)
        else:  # 系统消息转换为HumanMessage（临时处理方式）
            processed_messages.append(HumanMessage(content=msg))
    response = llm.invoke(processed_messages)
    return {"messages": messages + [response]}


# 1. 初始化图
graph = StateGraph(AdvancedChatState)

# 2. 添加所有节点
graph.add_node("judge_search", judge_search_node)  # 判断节点
graph.add_node("search", search_node)  # 搜索节点
graph.add_node("llm_answer", advanced_llm_node)  # LLM节点

# 3. 定义边：设置条件路由
# 入口 → 判断节点（无条件）
graph.set_entry_point("judge_search")

# 判断节点 → 分支（使用add_conditional_edges替代带condition的add_edge）
graph.add_conditional_edges(
    "judge_search",
    # 条件判断函数：返回下一个节点名
    lambda state: "search" if state["need_search"] else "llm_answer"
)

# 搜索节点 → LLM节点（搜索后，必须调用LLM生成回答）
graph.add_edge("search", "llm_answer")

# LLM节点 → 出口（生成回答后流程终止）
graph.set_finish_point("llm_answer")

# 4. 编译图并运行
compiled_graph = graph.compile()

# 测试1：需要搜索的问题
initial_state1 = {
    "messages": [HumanMessage(content="LangGraph 是什么？")],
    "need_search": False  # 初始值无关，会被判断节点更新
}
result1 = compiled_graph.invoke(initial_state1)

# 打印结果
print("测试1（需要搜索）：")
for msg in result1["messages"]:
    if isinstance(msg, (HumanMessage, AIMessage)):
        print(f"{msg.type.upper()}: {msg.content}\n")
    else:
        print(f"SYSTEM: {msg}\n")

# 测试2：不需要搜索的问题
initial_state2 = {
    "messages": [HumanMessage(content="用一句话介绍 LangGraph")],
    "need_search": False
}
result2 = compiled_graph.invoke(initial_state2)

print("测试2（不需要搜索）：")
for msg in result2["messages"]:
    if isinstance(msg, (HumanMessage, AIMessage)):
        print(f"{msg.type.upper()}: {msg.content}\n")
    else:
        print(f"SYSTEM: {msg}\n")
